#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
SIMION 命令行自动化示例

演示如何使用 Python 的 subprocess 模块调用 SIMION 命令行，
实现 Refine 和粒子飞行模拟的自动化。

注意：
    - Lua 脚本必须在 IOB 文件中配置，不能通过命令行参数指定
    - SIMION 没有 --lua-script= 参数
    - --lua "code" 用于执行 Lua 代码，不是脚本文件
"""

import subprocess
import os
from pathlib import Path


class SIMIONAutomation:
    """SIMION 命令行自动化工具"""
    
    def __init__(self, simion_exe=r"D:\SIMION-2024\simion.exe"):
        """
        初始化 SIMION 自动化工具
        
        参数:
            simion_exe: SIMION 可执行文件的完整路径
        """
        self.simion_exe = simion_exe
        
        # 验证 SIMION 是否存在
        if not os.path.exists(simion_exe):
            print(f"⚠️  警告: SIMION 未找到: {simion_exe}")
            print(f"   请修改路径或确保 SIMION 已安装")
    
    def refine(self, pa_file, num_threads=None, verbose=True):
        """
        Refine PA 文件（求解电场/磁场）
        
        参数:
            pa_file: PA 文件路径（.pa 或 .pa#）
            num_threads: CPU 线程数（None 使用默认值）
            verbose: 是否显示详细输出
        
        返回:
            bool: 成功返回 True，失败返回 False
        """
        # 构建命令
        cmd = [self.simion_exe, '--nogui', 'refine', pa_file]
        
        # 添加多线程选项
        if num_threads:
            cmd.insert(2, '--num-threads')
            cmd.insert(3, str(num_threads))
        
        if verbose:
            print(f"\n{'='*60}")
            print(f"执行 Refine: {pa_file}")
            print(f"命令: {' '.join(cmd)}")
            print(f"{'='*60}")
        
        # 执行命令
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        # 处理结果
        if result.returncode == 0:
            if verbose:
                print("✓ Refine 成功")
                if result.stdout:
                    print("\n--- 输出 ---")
                    print(result.stdout)
            return True
        else:
            if verbose:
                print("✗ Refine 失败")
                if result.stderr:
                    print("\n--- 错误 ---")
                    print(result.stderr)
            return False
    
    def fly(self, iob_file, noprompt=True, verbose=True):
        """
        运行粒子飞行模拟
        
        参数:
            iob_file: IOB 文件路径
            noprompt: 是否禁用交互提示（批处理时建议启用）
            verbose: 是否显示详细输出
        
        返回:
            tuple: (成功/失败, 输出文本)
        
        重要提示:
            Lua 脚本必须在 IOB 文件中配置！
            - 在 SIMION GUI 中: View → Workbench Program → 选择 .lua 文件
            - 不能通过命令行参数指定 Lua 脚本
            - SIMION 没有 --lua-script= 参数
        """
        # 构建命令
        cmd = [self.simion_exe, '--nogui']
        
        if noprompt:
            cmd.append('--noprompt')
        
        cmd.extend(['fly', iob_file])
        
        if verbose:
            print(f"\n{'='*60}")
            print(f"执行粒子飞行: {iob_file}")
            print(f"命令: {' '.join(cmd)}")
            print(f"{'='*60}")
        
        # 执行命令
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        # 处理结果
        if verbose:
            if result.returncode == 0:
                print("✓ 粒子飞行完成")
            else:
                print("✗ 粒子飞行失败")
            
            # 显示输出
            if result.stdout:
                print("\n--- 输出 ---")
                print(result.stdout)
            if result.stderr:
                print("\n--- 错误 ---")
                print(result.stderr)
        
        return (result.returncode == 0, result.stdout)
    
    def batch_refine(self, pa_dir, pattern='*.pa#', num_threads=None):
        """
        批量 Refine 多个 PA 文件
        
        参数:
            pa_dir: PA 文件所在目录
            pattern: 文件匹配模式（默认 *.pa#）
            num_threads: CPU 线程数
        
        返回:
            dict: {'成功': [...], '失败': [...]}
        """
        pa_path = Path(pa_dir)
        pa_files = list(pa_path.glob(pattern))
        
        results = {'成功': [], '失败': []}
        
        print(f"\n{'='*60}")
        print(f"批量 Refine")
        print(f"目录: {pa_dir}")
        print(f"模式: {pattern}")
        print(f"找到 {len(pa_files)} 个 PA 文件")
        print(f"{'='*60}")
        
        for i, pa_file in enumerate(pa_files, 1):
            print(f"\n[{i}/{len(pa_files)}] 处理 {pa_file.name}")
            
            if self.refine(str(pa_file), num_threads=num_threads, verbose=False):
                results['成功'].append(pa_file.name)
                print(f"  ✓ 成功")
            else:
                results['失败'].append(pa_file.name)
                print(f"  ✗ 失败")
        
        print(f"\n{'='*60}")
        print(f"完成: {len(results['成功'])} 成功, {len(results['失败'])} 失败")
        print(f"{'='*60}")
        
        return results
    
    def check_simion(self):
        """检查 SIMION 是否可用"""
        try:
            result = subprocess.run(
                [self.simion_exe, '--version'],
                capture_output=True,
                text=True,
                timeout=5
            )
            if result.returncode == 0:
                print(f"✓ SIMION 可用")
                print(result.stdout)
                return True
            else:
                print(f"✗ SIMION 不可用")
                return False
        except Exception as e:
            print(f"✗ SIMION 检查失败: {e}")
            return False


def parse_fly_output(output_text):
    """
    解析 SIMION 粒子飞行输出
    
    参数:
        output_text: SIMION 的标准输出文本
    
    返回:
        list: 粒子数据列表，每个元素为 dict，包含 x, y, z, velocity
    
    示例输出格式:
        status2,  1 Splat: xyz(     90,    171,   1399)mm, v=  1134.507mm/usec
        status2,  2 Splat: xyz(     84,    169,   1399)mm, v=  1134.607mm/usec
    """
    results = []
    
    for line in output_text.split('\n'):
        if 'Splat:' in line:
            try:
                # 提取粒子编号
                particle_num = int(line.split('Splat:')[0].strip().split()[-1])
                
                # 提取坐标: xyz(x, y, z)mm
                xyz_part = line.split('xyz(')[1].split(')mm')[0]
                x, y, z = map(float, [s.strip() for s in xyz_part.split(',')])
                
                # 提取速度: v=velocity mm/usec
                v_part = line.split('v=')[1].split('mm/usec')[0]
                velocity = float(v_part.strip())
                
                results.append({
                    'particle': particle_num,
                    'x': x,
                    'y': y,
                    'z': z,
                    'velocity': velocity
                })
            except Exception as e:
                print(f"⚠️  解析失败: {line}")
                print(f"   错误: {e}")
    
    return results


# ==================== 使用示例 ====================

def example_01_check_simion():
    """示例 1: 检查 SIMION 是否可用"""
    print("\n" + "="*60)
    print("示例 1: 检查 SIMION")
    print("="*60)
    
    auto = SIMIONAutomation(simion_exe=r"D:\SIMION-2024\simion.exe")
    auto.check_simion()


def example_02_refine_single_file():
    """示例 2: Refine 单个 PA 文件"""
    print("\n" + "="*60)
    print("示例 2: Refine 单个文件")
    print("="*60)
    
    auto = SIMIONAutomation(simion_exe=r"D:\SIMION-2024\simion.exe")
    
    # Refine 一个 PA 文件
    pa_file = 'pa/parallel_plates.pa#'
    
    if os.path.exists(pa_file):
        auto.refine(pa_file, num_threads=8)
    else:
        print(f"⚠️  文件不存在: {pa_file}")
        print(f"   请先运行 example_02_create_simple_electrode.py")


def example_03_fly_simulation():
    """示例 3: 运行粒子飞行模拟"""
    print("\n" + "="*60)
    print("示例 3: 粒子飞行模拟")
    print("="*60)
    
    auto = SIMIONAutomation(simion_exe=r"D:\SIMION-2024\simion.exe")
    
    # 运行粒子飞行（需要先准备好 IOB 文件）
    iob_file = '160x320.iob'
    
    if os.path.exists(iob_file):
        success, output = auto.fly(iob_file, noprompt=True)
        
        if success:
            # 解析输出
            particles = parse_fly_output(output)
            
            print(f"\n{'='*60}")
            print(f"检测到 {len(particles)} 个粒子")
            print(f"{'='*60}")
            
            for p in particles:
                print(f"粒子 {p['particle']}: "
                      f"位置=({p['x']:.1f}, {p['y']:.1f}, {p['z']:.1f}) mm, "
                      f"速度={p['velocity']:.2f} mm/µs")
    else:
        print(f"⚠️  文件不存在: {iob_file}")
        print(f"   请在 SIMION GUI 中创建 IOB 文件")
        print(f"   或使用您自己的 IOB 文件路径")


def example_04_batch_refine():
    """示例 4: 批量 Refine"""
    print("\n" + "="*60)
    print("示例 4: 批量 Refine")
    print("="*60)
    
    auto = SIMIONAutomation(simion_exe=r"D:\SIMION-2024\simion.exe")
    
    # 批量 Refine pa/ 目录下的所有 .pa# 文件
    pa_dir = 'pa/'
    
    if os.path.exists(pa_dir):
        results = auto.batch_refine(pa_dir, pattern='*.pa#', num_threads=8)
        
        print("\n成功的文件:")
        for file in results['成功']:
            print(f"  ✓ {file}")
        
        if results['失败']:
            print("\n失败的文件:")
            for file in results['失败']:
                print(f"  ✗ {file}")
    else:
        print(f"⚠️  目录不存在: {pa_dir}")
        print(f"   请先运行 example_02 或其他示例创建 PA 文件")


def example_05_parameter_sweep():
    """示例 5: 参数扫描（演示概念）"""
    print("\n" + "="*60)
    print("示例 5: 参数扫描（演示）")
    print("="*60)
    
    from SIMION.PA import PA
    
    auto = SIMIONAutomation(simion_exe=r"D:\SIMION-2024\simion.exe")
    
    # 创建多个不同电压的电极
    voltages = [100, 200, 300, 400, 500]
    
    print("\n创建 PA 文件:")
    for voltage in voltages:
        # 创建简单的平行板
        pa = PA(nx=100, ny=100, nz=1, symmetry='planar', max_voltage=1000)
        
        # 上板
        for x in range(100):
            pa.point(x, 20, 0, 1, voltage)
        
        # 下板
        for x in range(100):
            pa.point(x, 80, 0, 1, -voltage)
        
        # 保存
        pa_file = f'pa/sweep_v{voltage}.pa#'
        pa.save(pa_file)
        print(f"  ✓ 创建: {pa_file}")
    
    # 批量 Refine
    print("\n批量 Refine:")
    results = auto.batch_refine('pa/', pattern='sweep_*.pa#', num_threads=8)
    
    # 在实际应用中，这里会运行模拟并分析结果
    print(f"\n✓ 参数扫描完成")
    print(f"   生成了 {len(voltages)} 个不同电压的电极设计")
    print(f"   Refine 成功: {len(results['成功'])} 个")


def main():
    """主函数：运行所有示例"""
    print("\n" + "="*70)
    print("  SIMION Python 自动化示例")
    print("="*70)
    
    # 检查当前目录
    print(f"\n当前工作目录: {os.getcwd()}")
    
    # 运行示例
    try:
        example_01_check_simion()
        
        # 其他示例需要相应的文件存在
        # 取消注释来运行
        
        # example_02_refine_single_file()
        # example_03_fly_simulation()
        # example_04_batch_refine()
        # example_05_parameter_sweep()
        
    except Exception as e:
        print(f"\n✗ 错误: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n" + "="*70)
    print("提示:")
    print("  • 修改 simion_exe 路径以匹配您的 SIMION 安装位置")
    print("  • Lua 脚本必须在 IOB 文件中配置，不能通过命令行指定")
    print("  • 使用 --num-threads 可以加速 Refine 过程")
    print("  • 使用 --noprompt 可以避免交互提示（批处理）")
    print("="*70)


if __name__ == '__main__':
    main()

